#include "socks4a.h"

#include <memory>

#include "muduo/net/Endian.h"
using namespace muduo;
using namespace muduo::net;

Socks4a::Socks4a(EventLoop *loop, const InetAddress &listenAddr)
: TcpServer(loop, listenAddr, "Socks4")
{
  this->setConnectionCallback(std::bind(&Socks4a::onServerConnection, this, _1));
  this->setMessageCallback(std::bind(&Socks4a::onServerMessage, this, _1, _2, _3));
}
/**
 * @brief 处理收到客户端发来连接的情况
 *
 * @param conn
 */
void Socks4a::onServerConnection(const TcpConnectionPtr &conn)
{
  LOG_DEBUG << conn->name() << (conn->connected() ? " UP" : " DOWN");
  if (conn->connected()) {
    // 禁用Nagle算法，避免连续发包出现延迟，这对编写低延迟网络服务很重要
    conn->setTcpNoDelay(true);
  } else {
    auto it = tunnels_.find(conn->name());
    if (it != tunnels_.end()) {
      it->second->disconnect();
      tunnels_.erase(it);
    }
  }
}
/**
 * @brief 解析socks4协议，目前只支持41开头，即socks4*
 * connect，成功就返回地址对应的unique_ptr，否则返回nullptr
 *
 * @param data 需要解析的字符串
 * @param len  字符串的长度
 * @param where  是返回值，最后会指向域名的后一位，即socks4 中的 NULL，失败则指向nullptr
 * @return std::unique_ptr<InetAddress>
 */
static std::unique_ptr<InetAddress> codecSocks4(const char *data, size_t len, const char *&where)
{
  const char *begin = data + 8;
  const char *end = data + len;
  where = std::find(begin, end, '\0');
  if (where != end) {
    char ver = data[0];
    char cmd = data[1];
    const void *port = data + 2;
    const void *ip = data + 4;

    sockaddr_in addr;
    memZero(&addr, sizeof addr);
    addr.sin_family = AF_INET;
    addr.sin_port = *static_cast<const in_port_t *>(port);
    addr.sin_addr.s_addr = *static_cast<const uint32_t *>(ip);
    // 前24位都是0，那么就< 256,也就是socks4a，否则就是socks4
    bool socks4a = sockets::networkToHost32(addr.sin_addr.s_addr) < 256;
    bool okay = false;

    if (socks4a) {
      const char *endOfHostName = std::find(where + 1, end, '\0');

      if (endOfHostName != end) {
        std::string hostname = where + 1;
        where = endOfHostName;
        LOG_INFO << "Socks4a host name " << hostname;
        InetAddress tmp;
        // 获取域名对应的ip
        if (InetAddress::resolve(hostname, &tmp)) {
          addr.sin_addr.s_addr = tmp.ipv4NetEndian();
          okay = true;
        }
      } else
        return nullptr;  // 说明不是socks4*
    } else
      okay = true;

    // 是socks4*
    if (ver == 4 && cmd == 1 && okay) return std::make_unique<InetAddress>(addr);
  } else
    where = nullptr;

  return nullptr;
}
/**
 * @brief 处理客户端来的消息，以及第一个消息保存conn
 *
 * @param conn 与客户端的连接
 * @param buf 客户端发来的buf
 */
void Socks4a::onServerMessage(const TcpConnectionPtr &conn, Buffer *buf, muduo::Timestamp)
{
  LOG_DEBUG << conn->name() << " " << buf->readableBytes();
  if (tunnels_.find(conn->name()) == tunnels_.end()) {
    size_t len = buf->readableBytes();
    if (len > 128) {
      conn->shutdown();
    } else if (len > 8) {
      const char *where;
      auto serverAddr = codecSocks4(buf->peek(), len, where);
      // 不是socks4*
      if (serverAddr == nullptr) {
        // 找不到'\0'，收到了nullptr，说明不是socks4*
        if(where != nullptr) {
          char response[] = "\000\x5bUVWXYZ";
          conn->send(response, 8);
        }
        conn->shutdown();
      } else {
        // 开启通道
        auto tunnel = std::make_shared<Tunnel>(this->getLoop(), *serverAddr, conn);
        tunnel->setup();
        tunnel->connect();
        tunnels_[conn->name()] = tunnel;
        buf->retrieveUntil(where + 1);
        // 发送回应
        char response[] = "\000\x5aUVWXYZ";
        auto sin_port{serverAddr->portNetEndian()};
        auto s_addr{serverAddr->ipv4NetEndian()};
        memcpy(response + 2, &sin_port, 2);
        memcpy(response + 4, &s_addr, 4);
        conn->send(response, 8);
      }
    }
  } else if (!conn->getContext().empty()) {
    const TcpConnectionPtr &clientConn =
        boost::any_cast<const TcpConnectionPtr &>(conn->getContext());
    clientConn->send(buf);
  }
}
